package com.liuqi.aiagent.service.impl;

import com.liuqi.aiagent.mapper.TokensUsageMapper;
import com.liuqi.aiagent.service.TokensUsageService;
import com.liuqi.aiagent.transfer.ApiResult;
import com.liuqi.aiagent.transfer.dto.TokensUsageFindDto;
import com.liuqi.aiagent.transfer.entity.TokensUsage;
import com.liuqi.aiagent.transfer.model.Result;
import com.liuqi.aiagent.transfer.vo.TokensUsageStatisticsVo;
import com.liuqi.aiagent.util.DateUtil;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneId;
import java.time.temporal.ChronoUnit;
import java.util.ArrayList;
import java.util.List;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

/**
 * @author liuqi
 * @date 2025/7/24
 **/
@Service
@Slf4j
@AllArgsConstructor
public class TokensUsageServiceImpl implements TokensUsageService {

    private TokensUsageMapper tokensUsageMapper;

    @Override
    public Result<TokensUsageStatisticsVo> statistics(TokensUsageFindDto findDto) {
        String begin = findDto.getBeginDate();
        String end = findDto.getEndDate();

        findDto.setBeginDate(begin.concat(" 00:00:00"));
        findDto.setEndDate(end.concat(" 23:59:59"));
        List<TokensUsage> tokensUsages = tokensUsageMapper.find(findDto);

        TokensUsageStatisticsVo statisticsVo = new Aggregation(tokensUsages, begin, end, findDto.getDateType()).aggregation();

        return ApiResult.success(statisticsVo);
    }

    static class Aggregation {
        final List<TokensUsage> tokensUsages;
        final String beginDate;
        final String endDate;
        final TokensUsageFindDto.DateType dateType;
        final int window;
        final TokensUsageStatisticsVo statisticsVo = new TokensUsageStatisticsVo();

        int currentIndex = 0;

        long currentTimestamp;
        long nextTimestamp;

        Aggregation(List<TokensUsage> tokensUsages, String beginDate, String endDate, TokensUsageFindDto.DateType dateType) {
            this.tokensUsages = tokensUsages;
            this.beginDate = beginDate;
            this.endDate = endDate;
            this.dateType = dateType;

            LocalDate begin = LocalDate.parse(beginDate, DateUtil.Y_M_D);
            LocalDate end = LocalDate.parse(endDate, DateUtil.Y_M_D);

            setTimestamp(begin);

            if (dateType == TokensUsageFindDto.DateType.MONTH) {
                this.window = (int) ChronoUnit.MONTHS.between(begin, end) + 1;
            } else {
                this.window = (int) ChronoUnit.DAYS.between(begin, end) + 1;
            }

            statisticsVo.setXData(new ArrayList<>(this.window));
            statisticsVo.setInput(new ArrayList<>(this.window));
            statisticsVo.setOutput(new ArrayList<>(this.window));
            statisticsVo.setTotal(new ArrayList<>(this.window));
        }

        TokensUsageStatisticsVo aggregation() {
            for (int i = 0; i < this.window; i++) {
                statisticsVo.getXData().add(formX());

                long input = 0;
                long output = 0;
                long total = 0;

                while (currentIndex < tokensUsages.size()) {
                    TokensUsage tokensUsage = tokensUsages.get(currentIndex);
                    long time = tokensUsage.getCreateTime().getTime();
                    if (time <= nextTimestamp) {
                        input += tokensUsage.getInputToken();
                        output += tokensUsage.getOutputToken();
                        total += tokensUsage.getTotalToken();
                        currentIndex++;
                    } else {
                        break;
                    }
                }

                statisticsVo.getInput().add(input);
                statisticsVo.getOutput().add(output);
                statisticsVo.getTotal().add(total);

                nextTimestamp();
            }

            return statisticsVo;
        }

        void setTimestamp(LocalDate localDate) {
            if (dateType == TokensUsageFindDto.DateType.DAY) {
                this.currentTimestamp = DateUtil.startDayTimestamp(localDate);
                this.nextTimestamp = DateUtil.endDayTimestamp(localDate);
            } else {
                this.currentTimestamp = DateUtil.startMonthTimestamp(localDate);
                this.nextTimestamp = DateUtil.endMonthTimestamp(localDate);
            }
        }

        void nextTimestamp() {
            LocalDate localDate = Instant.ofEpochMilli(currentTimestamp).atZone(ZoneId.systemDefault()).toLocalDate();
            if (dateType == TokensUsageFindDto.DateType.DAY) {
                localDate = localDate.plusDays(1);
            } else {
                localDate = localDate.plusMonths(1);
            }
            setTimestamp(localDate);
        }

        String formX() {
            if (dateType == TokensUsageFindDto.DateType.DAY) {
                return Instant.ofEpochMilli(currentTimestamp).atZone(ZoneId.systemDefault()).toLocalDate().format(DateUtil.Y_M_D);
            } else {
                return Instant.ofEpochMilli(currentTimestamp).atZone(ZoneId.systemDefault()).toLocalDate().format(DateUtil.Y_M);
            }
        }

    }

}
